#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import argparse
import csv
import json
import re
from pathlib import Path
from typing import Dict, Any, List

def normalize_key_from_image(image: str) -> str:
    """
    将 JSON 中的 image 字段规范化为 CSV 中的 Source 形式:
    '/.../dataset/<uuid>.pdf::page=0' -> './dataset/<uuid>.pdf'
    '/.../dataset/<uuid>.png'         -> './dataset/<uuid>.png'
    """
    m = re.search(r"(dataset/[^:]+)", image)
    if not m:
        return ""
    return "./" + m.group(1)

def top1_caption(item: Dict[str, Any]) -> str:
    """
    取 topn 的第一条 caption，优先 rank==1；否则取列表第一条
    """
    topn: List[Dict[str, Any]] = item.get("topn", []) or []
    if not topn:
        return ""
    # 优先 rank==1
    for t in topn:
        if t.get("rank") == 1 and isinstance(t.get("caption"), str):
            return t["caption"]
    # 否则第一条
    first = topn[0].get("caption")
    return first if isinstance(first, str) else ""

def build_mapping(json_path: Path) -> Dict[str, str]:
    data = json.loads(json_path.read_text(encoding="utf-8"))
    mapping: Dict[str, str] = {}
    if not isinstance(data, list):
        raise ValueError("JSON 顶层应为数组（list）。")
    for item in data:
        img = item.get("image", "")
        key = normalize_key_from_image(img)
        if not key:
            continue
        cap = top1_caption(item)
        if cap:
            # 若同一 key 出现多次，保留第一个（或按需覆盖也可）
            mapping.setdefault(key, cap)
    return mapping

def update_csv(csv_in: Path, csv_out: Path, mapping: Dict[str, str]) -> None:
    # 读取并保持原顺序；只改 Caption 列
    with csv_in.open("r", encoding="utf-8-sig", newline="") as f:
        reader = list(csv.DictReader(f))
        if not reader:
            # 空文件则原样写出表头
            with csv_out.open("w", encoding="utf-8", newline="") as wf:
                writer = csv.DictWriter(wf, fieldnames=["Source", "Caption"])
                writer.writeheader()
            return

        # 校验列
        fieldnames = reader[0].keys()
        if "Source" not in fieldnames or "Caption" not in fieldnames:
            raise ValueError("CSV 必须包含列：Source, Caption")

        updated_rows = []
        for row in reader:
            src = row.get("Source", "")
            if src in mapping:
                row["Caption"] = mapping[src]
            # 否则保留原 Caption
            updated_rows.append(row)

    with csv_out.open("w", encoding="utf-8", newline="") as wf:
        writer = csv.DictWriter(wf, fieldnames=["Source", "Caption"])
        writer.writeheader()
        for row in updated_rows:
            # 仅写出需要的两列，防止多余列顺带写出
            writer.writerow({"Source": row["Source"], "Caption": row["Caption"]})

def main():
    ap = argparse.ArgumentParser(
        description="将 JSON 中 image->top1 caption 映射到 CSV（保持 Source 与顺序不变，仅更新 Caption）"
    )
    ap.add_argument("--json", required=True, type=Path, help="输入 JSON 路径")
    ap.add_argument("--csv", required=True, type=Path, help="输入 CSV 路径（含列：Source, Caption）")
    ap.add_argument("--out", required=True, type=Path, help="输出 CSV 路径")
    args = ap.parse_args()

    mapping = build_mapping(args.json)
    update_csv(args.csv, args.out, mapping)

if __name__ == "__main__":
    main()
